/*
Copyright (C) 1996-1997 Id Software, Inc.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 3
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

*/

// i had to switch precompiled headers off for this file because D3Dcompiler.h now has dependencies on d3d11.h (thanks Bill)
// but d3d11.h can't be included after d3d10.h (thanks again Bill)
#include <D3D11.h>
#include <D3Dcompiler.h>

#include "quakedef.h"
#include <typeinfo>  //for 'typeid' to work

// needed for the shader compiler
#pragma comment (lib, "d3dcompiler.lib")

// shortcuts for defining registers so that I don't need to mess with crap
#define DEFINE_TEXTURE_REGISTER(n) {#n, va ("t%i", n)}
#define DEFINE_BUFFER_REGISTER(n) {#n, va ("b%i", n)}
#define DEFINE_SAMPLER_REGISTER(n) {#n, va ("s%i", n)}

// all shader types inherit these defines
#define SHARED_DEFINES \
	{"MAX_LIGHTSTYLES", va ("%i", MAX_LIGHTSTYLES)}, \
	{"MAX_PACKED_LIGHTSTYLES", va ("%i", MAX_LIGHTSTYLES >> 2)}, \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_SPRITES), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_SURFACES), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_SOLIDSKY), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_ALPHASKY), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_CUBESKY), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_LIGHTMAP0), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_LIGHTMAP1), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_LIGHTMAP2), \
	DEFINE_TEXTURE_REGISTER (TEXTURE_SLOT_LIGHTMAP3)

// common input layout definitions
static const D3D10_INPUT_ELEMENT_DESC SurfLayout[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("TEXCOORD", 0, DXGI_FORMAT_R16G16_FLOAT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("LIGHTMAP", 0, DXGI_FORMAT_R16G16_FLOAT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("MAPNUM", 0, DXGI_FORMAT_R32_UINT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("STYLES", 0, DXGI_FORMAT_R8G8B8A8_UINT, BUFFER_SLOT_SURFACES, 0)
};

static const D3D10_INPUT_ELEMENT_DESC MeshLayout[] = {
	VDECL ("PREVPOSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 0),
	VDECL ("PREVNORMAL", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 0),
	VDECL ("CURRPOSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 1, 0),
	VDECL ("CURRNORMAL", 0, DXGI_FORMAT_R32G32B32_FLOAT, 1, 0),
	VDECL ("TEXCOORD", 0, DXGI_FORMAT_R16G16_FLOAT, 2, 0)
};

static const D3D10_INPUT_ELEMENT_DESC PositionXYZ[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_LVOLUME, 0)
};

static const D3D10_INPUT_ELEMENT_DESC WarpLayout[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("TEXCOORD", 0, DXGI_FORMAT_R16G16_FLOAT, BUFFER_SLOT_SURFACES, 0),
	VDECL ("WARPCOORD", 0, DXGI_FORMAT_R16G16_FLOAT, BUFFER_SLOT_SURFACES, 0)
};

static const D3D10_INPUT_ELEMENT_DESC SpriteLayout[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32_FLOAT, BUFFER_SLOT_SPRITES, 0),
	VDECL ("TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, BUFFER_SLOT_SPRITES, 0)
};

static const D3D10_INPUT_ELEMENT_DESC DrawLayout[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, BUFFER_SLOT_DRAW, 1),
	VDECL ("TEXCOORD", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, BUFFER_SLOT_DRAW, 1),
	VDECL ("COLOUR", 0, DXGI_FORMAT_R8G8B8A8_UNORM, BUFFER_SLOT_DRAW, 1)
};

static const D3D10_INPUT_ELEMENT_DESC PartLayout[] = {
	VDECL ("ORIGIN", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_PARTICLES, 0),
	VDECL ("COLORLO", 0, DXGI_FORMAT_R8G8B8A8_UNORM, BUFFER_SLOT_PARTICLES, 0),
	VDECL ("COLORHI", 0, DXGI_FORMAT_R8G8B8A8_UNORM, BUFFER_SLOT_PARTICLES, 0),
	VDECL ("COLORTIME", 0, DXGI_FORMAT_R32_FLOAT, BUFFER_SLOT_PARTICLES, 0)
};

static const D3D10_INPUT_ELEMENT_DESC LightLayout[] = {
	VDECL ("ORIGIN", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_LIGHTS, 1),
	VDECL ("RADIUS", 0, DXGI_FORMAT_R32_FLOAT, BUFFER_SLOT_LIGHTS, 1),
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_LVOLUME, 0)
};

static const D3D10_INPUT_ELEMENT_DESC SkyLayout[] = {
	VDECL ("POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, BUFFER_SLOT_SURFACES, 0)
};

static const D3D10_INPUT_ELEMENT_DESC MeshLayoutSV[] = {
	VDECL ("PREVPOSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0, 0),
	VDECL ("CURRPOSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 1, 0)
};

// gather together a layout and it's elements
struct QINPUTLAYOUT
{
	const D3D10_INPUT_ELEMENT_DESC *Elements;
	int NumElements;
};


// save me typing
#define DEFINE_LAYOUT(layout) {layout, sizeof (layout) / sizeof (layout[0])}

// and define them all (these must match the definition order in glquake.h)
static const QINPUTLAYOUT d3d_InputLayouts[] = {
	DEFINE_LAYOUT (SurfLayout),
	DEFINE_LAYOUT (MeshLayout),
	DEFINE_LAYOUT (PositionXYZ),
	DEFINE_LAYOUT (WarpLayout),
	DEFINE_LAYOUT (SpriteLayout),
	DEFINE_LAYOUT (DrawLayout),
	DEFINE_LAYOUT (PartLayout),
	DEFINE_LAYOUT (LightLayout),
	DEFINE_LAYOUT (SkyLayout),
	DEFINE_LAYOUT (MeshLayoutSV)
};


struct shadercache_t
{
	char EntryPoint[128];
	char ShaderType[128];
	unsigned short CRC0;
	unsigned short CRC1;
	ID3D10DeviceChild *Shader;
};


#define MAX_SHADER_CACHE	1024

shadercache_t ShaderCache[MAX_SHADER_CACHE];
int NumShaderCache = 0;

class QShaderIncludes : public ID3D10Include
{
public:
	HRESULT WINAPI Open (D3D10_INCLUDE_TYPE IncludeType, LPCSTR pFileName, LPCVOID pParentData, LPCVOID *ppData, UINT *pBytes)
	{
		char *included = NULL;
		int includelen = 0;

		// check for specific files we wish to include (file names aren't stored in the exe - boo)
		if (!strcmp (pFileName, "fxVSIncludes.fx"))
			includelen = Sys_LoadResourceData (IDR_VSINCLUDES, (void **) &included);

		if (includelen && included)
		{
			*ppData = included;
			*pBytes = includelen;
		}
		else
		{
			*ppData = NULL;
			*pBytes = 0;
		}

		return S_OK;
	}

	HRESULT WINAPI Close (LPCVOID pData)
	{
		return S_OK;
	}
};


int lastresourceid = -1;
const char *lastresourcetext = NULL;
int lastresourcelength = 0;

class ShaderFactory
{
public:
	ShaderFactory (int resourceID)
	{
		if (resourceID == lastresourceid)
		{
			this->textlength = lastresourcelength;
			this->shadertext = lastresourcetext;
		}
		else
		{
			// load the shader resource text
			this->textlength = Sys_LoadResourceData (resourceID, (void **) &this->shadertext);
			lastresourcetext = this->shadertext;
			lastresourcelength = this->textlength;
		}

		// set up everything else
		this->entrypoint = NULL;
		this->shadertype = NULL;
		this->profile = NULL;
		this->defines = NULL;
		this->ShaderBlob = NULL;
		this->ErrorBlob = NULL;

		lastresourceid = resourceID;
	}

	~ShaderFactory ()
	{
		// release blob objects
		SAFE_RELEASE (this->ShaderBlob);
		SAFE_RELEASE (this->ErrorBlob);
	}

	const char *shadertext;
	int textlength;
	const char *entrypoint;
	const char *shadertype;
	const char *profile;
	const D3D10_SHADER_MACRO *defines;
	ID3D10Blob *ErrorBlob;
	ID3D10Blob *ShaderBlob;

private:
	unsigned short CRC0;
	unsigned short CRC1;

	unsigned short GetCRC (void)
	{
		return CRC_Block (this->ShaderBlob->GetBufferPointer (), this->ShaderBlob->GetBufferSize ());
	}

	template <typename ShaderType>
	BOOL CheckCache (ShaderType **Shader, const D3D10_INPUT_ELEMENT_DESC *elements = NULL, int numelements = 0)
	{
		this->shadertype = typeid (Shader).name ();

		// CRC the buffer
		this->CRC0 = this->GetCRC ();
		this->CRC1 = 0;

		// optionally CRC the input layout definition
		if (elements && numelements)
			this->CRC1 = CRC_Block (elements, sizeof (D3D10_INPUT_ELEMENT_DESC) * numelements);

		// look for it in the cache
		for (int i = 0; i < NumShaderCache; i++)
		{
			// check fast matches first
			if (!ShaderCache[i].Shader) continue;
			if (ShaderCache[i].CRC0 != this->CRC0) continue;
			if (ShaderCache[i].CRC1 != this->CRC1) continue;

			// then check slow matches
			if (strcmp (ShaderCache[i].ShaderType, this->shadertype)) continue;
			if (strcmp (ShaderCache[i].EntryPoint, this->entrypoint)) continue;

			Con_Printf (PRINT_DEVELOPER, "reusing %s %s from shader cache\n", this->shadertype, ShaderCache[i].EntryPoint);

			// this shader matches one in the cache so hand that back
			*Shader = (ShaderType *) ShaderCache[i].Shader;
			return TRUE;
		}

		return FALSE;
	}

	template <typename ShaderType>
	void CacheShader (ShaderType **Shader, const D3D10_INPUT_ELEMENT_DESC *elements = NULL, int numelements = 0)
	{
		if (NumShaderCache >= MAX_SHADER_CACHE)
			Sys_Error ("R_CacheShader : NumShaderCache >= MAX_SHADER_CACHE");

		// cache this shader
		shadercache_t *sc = &ShaderCache[NumShaderCache];
		NumShaderCache++;

		// copy over the shader
		sc->Shader = *Shader;

		// set the name and type
		strcpy (sc->EntryPoint, this->entrypoint);
		strcpy (sc->ShaderType, this->shadertype);

		// copy over the rest of the data
		sc->CRC0 = this->CRC0;
		sc->CRC1 = this->CRC1;
	}

	void ShaderError (HRESULT hr, char *errmsg)
	{
		if (SUCCEEDED (hr))
			Con_Printf (PRINT_SAFE, "%s\n", errmsg);
		else Sys_Error (errmsg);
	}

	HRESULT CompileShader (const char *entrypoint, const char *profile, const D3D10_SHADER_MACRO *defines)
	{
		// release blob objects
		SAFE_RELEASE (this->ShaderBlob);
		SAFE_RELEASE (this->ErrorBlob);

		this->entrypoint = entrypoint;
		this->profile = profile;
		this->defines = defines;

		QShaderIncludes Includes;

		HRESULT hr = D3DCompile (
			this->shadertext,
			this->textlength,
			NULL,
			this->defines,
			&Includes,
			this->entrypoint,
			this->profile,
			0,
			0,
			&this->ShaderBlob,
			&this->ErrorBlob
		);

		if (this->ErrorBlob != NULL)
		{
			this->ShaderError (hr, va ("Error compiling %s\n%s", this->entrypoint, (char *) this->ErrorBlob->GetBufferPointer ()));
			this->ErrorBlob->Release ();
			this->ErrorBlob = NULL;
		}
		else if (FAILED (hr))
			this->ShaderError (hr, "Terminal shader error!!!");

		return hr;
	}

	void CheckCreate (HRESULT hr)
	{
		if (FAILED (hr))
		{
			Sys_Error ("FAILED (hr) creating %s\nCompilation succeeded but shader creation failed", this->entrypoint);
		}
	}

public:
	void CreateVertexShaderAndInputLayout (ID3D10VertexShader **VS, ID3D10InputLayout **IL, const char *vsentry, int layout)
	{
		const D3D10_SHADER_MACRO defines[] = {
			{"VERTEXSHADER", "1"},
			SHARED_DEFINES,
			{NULL, NULL}
		};

		if (SUCCEEDED (this->CompileShader (vsentry, "vs_4_0", defines)))
		{
			if (!this->CheckCache (VS))
			{
				this->CheckCreate (d3d_Device->CreateVertexShader (
					this->ShaderBlob->GetBufferPointer (),
					this->ShaderBlob->GetBufferSize (),
					VS
				));

				this->CacheShader (VS);

				// create the shader without an associated input layout
				if (layout == INPUT_LAYOUT_NONE) return;

				// must do the layout here too otherwise the blobs won't be valid to use
				if (!this->CheckCache (IL, d3d_InputLayouts[layout].Elements, d3d_InputLayouts[layout].NumElements))
				{
					this->CheckCreate (d3d_Device->CreateInputLayout (
						d3d_InputLayouts[layout].Elements,
						d3d_InputLayouts[layout].NumElements,
						this->ShaderBlob->GetBufferPointer (),
						this->ShaderBlob->GetBufferSize (),
						IL
					));

					this->CacheShader (IL);
				}
			}
			else if (layout == INPUT_LAYOUT_NONE)
				return;
			else if (!this->CheckCache (IL, d3d_InputLayouts[layout].Elements, d3d_InputLayouts[layout].NumElements))
			{
				// this is a terminal error !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
				Sys_Error ("ShaderFactory::CreateVertexShaderAndInputLayout - vertex shader was cached but input layout was not");
			}
		}
	}

	void CreateGeometryShader (ID3D10GeometryShader **GS, const char *gsentry)
	{
		const D3D10_SHADER_MACRO defines[] = {
			{"GEOMETRYSHADER", "1"},
			SHARED_DEFINES,
			{NULL, NULL}
		};

		if (SUCCEEDED (this->CompileShader (gsentry, "gs_4_0", defines)))
		{
			if (!this->CheckCache (GS))
			{
				this->CheckCreate (d3d_Device->CreateGeometryShader (
					this->ShaderBlob->GetBufferPointer (),
					this->ShaderBlob->GetBufferSize (),
					GS
				));

				this->CacheShader (GS);
			}
		}
	}

	void CreatePixelShader (ID3D10PixelShader **PS, const char *psentry)
	{
		const D3D10_SHADER_MACRO defines[] = {
			{"PIXELSHADER", "1"},
			SHARED_DEFINES,
			{NULL, NULL}
		};

		if (SUCCEEDED (this->CompileShader (psentry, "ps_4_0", defines)))
		{
			if (!this->CheckCache (PS))
			{
				this->CheckCreate (d3d_Device->CreatePixelShader (
					this->ShaderBlob->GetBufferPointer (),
					this->ShaderBlob->GetBufferSize (),
					PS
				));

				this->CacheShader (PS);
			}
		}
	}
};


void R_FreeShaderCache (void)
{
	for (int i = 0; i < MAX_SHADER_CACHE; i++)
	{
		SAFE_RELEASE (ShaderCache[i].Shader);
		memset (&ShaderCache[i], 0, sizeof (ShaderCache[i]));
	}
}


QSHADER::QSHADER (void)
{
	this->VertexShader = NULL;
	this->InputLayout = NULL;
	this->GeometryShader = NULL;
	this->PixelShader = NULL;
}


void QSHADER::CreateShaders (int resourceID, const char *vsentry, const char *gsentry, const char *psentry, int layout)
{
	// because register defines can use a lot of scratch, we take a mark now and free it when done
	int mark = qScratch->GetLowMark ();

	ShaderFactory CreateStruct (resourceID);

	if (vsentry) CreateStruct.CreateVertexShaderAndInputLayout (&this->VertexShader, &this->InputLayout, vsentry, layout);
	if (gsentry) CreateStruct.CreateGeometryShader (&this->GeometryShader, gsentry);
	if (psentry) CreateStruct.CreatePixelShader (&this->PixelShader, psentry);

	qScratch->FreeToLowMark (mark);
}


void QSHADER::Bind (void)
{
	// always bind even if NULL so that we can deactivate a shader stage (e.g. GS)
	d3d_State->VSSetShader (this->VertexShader);
	d3d_State->IASetLayout (this->InputLayout);
	d3d_State->GSSetShader (this->GeometryShader);
	d3d_State->PSSetShader (this->PixelShader);
}

